import os
import cv2
from PIL import Image
import numpy as np

file_dir = "C:\\Users\\jiangshan\\Desktop\\Lock3DFace-demo\\color"
img_list = ["test.jpg"]
# for root, dirs, files in os.walk(file_dir):
#     if files:
#         for file in files:
#             img_path = root
#             img_list.append(root + "\\" + file)
#             img_path = img_path.replace("color", "rgbd")
#             img_path = img_path.replace("COLOR", "RGBD")
#             # print(img_path)
#             # print(os.path.exists(img_path))
#             if not os.path.exists(img_path):
#                 os.makedirs(img_path)

for color_path in img_list:
    depth_path = color_path.replace("test", "depth")
    # depth_path = depth_path.replace("COLOR", "DEPTH")
    depth_path = depth_path.replace("jpg", "png")

    rgb = cv2.imread(color_path)
    depth_img = Image.open(depth_path)
    depth = np.array(depth_img)

    rgbd = np.zeros((424, 512, 4), dtype=np.uint8)
    rgbd[:, :, 0] = rgb[:, :, 0]
    rgbd[:, :, 1] = rgb[:, :, 1]
    rgbd[:, :, 2] = rgb[:, :, 2]
    rgbd[:, :, 3] = depth

    # print(rgb.shape)
    # print(depth.shape)

    new_rgb = rgbd[:, :, 0:3]
    new_depth = rgbd[3]

    depth_color_image = cv2.convertScaleAbs(depth, alpha=256/8000)
    depth_color_image = cv2.applyColorMap(depth_color_image, cv2.COLORMAP_JET)

    # depth_color_image = cv2.applyColorMap(np.round(depth / 30).astype(np.uint8),
    #                                                   cv2.COLORMAP_JET)
    combined_image = cv2.addWeighted(new_rgb, 0.7, depth_color_image, 0.3, 0)

    cv2.imshow("combined", combined_image)
    cv2.waitKey(0)

    # rgbd_path = depth_path.replace("depth", "rgbd")
    # rgbd_path = rgbd_path.replace("DEPTH", "RGBD")
    #
    # cv2.imwrite(rgbd_path, rgbd)